package 中等.树;

import util.TreeNode;

import java.util.ArrayDeque;
import java.util.Deque;

/**
 * 给你二叉树的根结点 root ，此外树的每个结点的值要么是 0 ，要么是 1 。
 * <p>
 * 返回移除了所有不包含 1 的子树的原二叉树。
 * <p>
 * 节点 node 的子树为 node 本身加上所有 node 的后代。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode-cn.com/problems/binary-tree-pruning
 */
public class 二叉树剪枝_814 {

    public static void main(String[] args) {

        TreeNode root = new TreeNode(1);
        TreeNode node1 = new TreeNode(0);
        TreeNode node2 = new TreeNode(0);
        TreeNode node3 = new TreeNode(1);
        root.right = node1;
        node1.left = node2;
        node1.right = node3;

        pruneTree(root);

    }

    /**
     * 后序遍历
     * 叶子节点的值为 0 时，将节点置为 null
     * 如果该节点不是叶子节点，说明它的子树中存在 1
     *
     * @param root
     * @return
     */
    public static TreeNode pruneTree3(TreeNode root) {
        if (root == null) return null;
        root.left = pruneTree3(root.left);
        root.right = pruneTree3(root.right);
        if (root.val == 0 && root.left == null && root.right == null) return null;
        return root;
    }

    /**
     * 广度搜索遍历+递归查找左右子树是否包含1
     *
     * @param root
     * @return
     */
    public static TreeNode pruneTree(TreeNode root) {
        if (!hasOne(root)) {
            return null;
        }
        Deque<TreeNode> queue = new ArrayDeque<>();
        queue.addLast(root);
        while (!queue.isEmpty()) {
            int size = queue.size();
            while (size-- > 0) {
                TreeNode curNode = queue.pollFirst();
                if (!hasOne(curNode.left)) {
                    curNode.left = null;
                }
                if (!hasOne(curNode.right)) {
                    curNode.right = null;
                }

                if (curNode.right != null) {
                    queue.addLast(curNode.right);
                }
                if (curNode.left != null) {
                    queue.addLast(curNode.left);
                }
            }
        }
        return root;
    }

    public static boolean hasOne(TreeNode node) {
        if (node == null) {
            return false;
        }
        if (node.val == 1) {
            return true;
        }
        return hasOne(node.right) || hasOne(node.left);
    }

    /**
     * 在递归的过程中判断如果当前节点和其所有子节点不包含1，那么去除这个节点
     * 由于我们需要父节点才能将当前节点置空，那么递归中只判断，除了根节点的节点
     *
     * @param root
     * @return
     */
    public static TreeNode pruneTree2(TreeNode root) {
        return containsOne(root) ? root : null;
    }

    public static boolean containsOne(TreeNode node) {
        if (node == null) {
            return false;
        }
        boolean leftHasOne = containsOne(node.left);
        boolean rightHasOne = containsOne(node.right);
        if (!leftHasOne) {
            node.left = null;
        }
        if (!rightHasOne) {
            node.right = null;
        }
        return node.val == 1 || leftHasOne || rightHasOne;
    }

}
